#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/10/18 23:20
# @Author  : when
# @File    : sqlalchemy_connector.py
# @Description : SQLAlchemy数据连接器

from datetime import datetime

from sqlalchemy import create_engine
from sqlalchemy.orm import Session

try:
    from components.mysql_client.config import ENGINE_URL
except ImportError:
    ENGINE_URL = "mysql+pymysql://scott:tiger@hostname/dbname"

# engine = create_engine(ENGINE_URL, echo=True)
engine = create_engine(ENGINE_URL)


class SqlConnector:
    def __init__(self):
        self.session = Session(engine)

    def insert_one(self, obj, auto_commit=True):
        if hasattr(obj, "update_time"):
            obj.update_time = datetime.now()
        if hasattr(obj, "created_time"):
            obj.created_time = datetime.now()
        self.session.add(obj)
        self.session.flush()
        if not auto_commit:
            return
        self.commit()

    def insert(self, obj_list: list, auto_commit=True):
        if not isinstance(obj_list, list):
            obj_list = list(obj_list)
        for obj in obj_list:
            if hasattr(obj, "update_time"):
                obj.update_time = datetime.now()
            if hasattr(obj, "created_time"):
                obj.created_time = datetime.now()
        self.session.add_all(obj_list)
        self.session.flush()
        if not auto_commit:
            return
        self.commit()

    def query(self, obj_base):
        if hasattr(obj_base, "update_time"):
            obj_base.update_time = datetime.now()
        return self.session.query(obj_base).filter_by(is_removed="f")

    def update(self, obj_base):
        return self.session.query(obj_base).filter_by(is_removed="f")

    def commit(self):
        self.session.commit()

    def rollback(self):
        self.session.rollback()
